#include "postgres.h"

#include "plpy_elog.h"
#include "plpy_typeio.h"
#include "plpython.h"
#include "utils/fmgrprotos.h"
#include "utils/jsonb.h"
#include "utils/numeric.h"

PG_MODULE_MAGIC;

void		_PG_init(void);

/* for PLyObject_AsString in plpy_typeio.c */
typedef char *(*PLyObject_AsString_t) (PyObject *plrv);
static PLyObject_AsString_t PLyObject_AsString_p;

typedef void (*PLy_elog_impl_t) (int elevel, const char *fmt,...);
static PLy_elog_impl_t PLy_elog_impl_p;

/*
 * decimal_constructor 是来自 Python 库的一个函数，用于将字符串转换为 Python decimal 类型
 */
static PyObject *decimal_constructor;

static PyObject *fc_PLyObject_FromJsonbContainer(JsonbContainer *fc_jsonb);
static JsonbValue *fc_PLyObject_ToJsonbValue(PyObject *fc_obj,
										  JsonbParseState **fc_jsonb_state, bool fc_is_elem);

typedef PyObject *(*PLyUnicode_FromStringAndSize_t)
			(const char *s, Py_ssize_t size);
static PLyUnicode_FromStringAndSize_t PLyUnicode_FromStringAndSize_p;

/*
 * 模块初始化函数：获取跨模块调用的函数指针。
 */
void _PG_init(void)
{
	/* 断言验证上述 typedef 与原始声明匹配 */
	AssertVariableIsOfType(&PLyObject_AsString, PLyObject_AsString_t);
	PLyObject_AsString_p = (PLyObject_AsString_t)
		load_external_function("$libdir/" PLPYTHON_LIBNAME, "PLyObject_AsString",
							   true, NULL);
	AssertVariableIsOfType(&PLyUnicode_FromStringAndSize, PLyUnicode_FromStringAndSize_t);
	PLyUnicode_FromStringAndSize_p = (PLyUnicode_FromStringAndSize_t)
		load_external_function("$libdir/" PLPYTHON_LIBNAME, "PLyUnicode_FromStringAndSize",
							   true, NULL);
	AssertVariableIsOfType(&PLy_elog_impl, PLy_elog_impl_t);
	PLy_elog_impl_p = (PLy_elog_impl_t)
		load_external_function("$libdir/" PLPYTHON_LIBNAME, "PLy_elog_impl",
							   true, NULL);
}

/* 这些定义必须在 _PG_init 之后 */
#define PLyObject_AsString (PLyObject_AsString_p)
#define PLyUnicode_FromStringAndSize (PLyUnicode_FromStringAndSize_p)
#undef PLy_elog
#define PLy_elog (PLy_elog_impl_p)

/*
 * PLyUnicode_FromJsonbValue
 *
 * 将字符串 JsonbValue 转换为 Python 字符串。
 */
static PyObject * fc_PLyUnicode_FromJsonbValue(JsonbValue *fc_jbv)
{
	Assert(fc_jbv->type == jbvString);

	return PLyUnicode_FromStringAndSize(fc_jbv->val.string.val, fc_jbv->val.string.len);
}

/*
 * PLyUnicode_ToJsonbValue
 *
 * 将 Python 字符串转换为 JsonbValue。
 */
static void fc_PLyUnicode_ToJsonbValue(PyObject *fc_obj, JsonbValue *fc_jbvElem)
{
	fc_jbvElem->type = jbvString;
	fc_jbvElem->val.string.val = PLyObject_AsString(fc_obj);
	fc_jbvElem->val.string.len = strlen(fc_jbvElem->val.string.val);
}

/*
 * PLyObject_FromJsonbValue
 *
 * 将 JsonbValue 转换为 PyObject。
 */
static PyObject * fc_PLyObject_FromJsonbValue(JsonbValue *fc_jsonbValue)
{
	switch (fc_jsonbValue->type)
	{
		case jbvNull:
			Py_RETURN_NONE;

		case jbvBinary:
			return fc_PLyObject_FromJsonbContainer(fc_jsonbValue->val.binary.data);

		case jbvNumeric:
			{
				Datum		fc_num;
				char	   *fc_str;

				fc_num = NumericGetDatum(fc_jsonbValue->val.numeric);
				fc_str = DatumGetCString(DirectFunctionCall1(numeric_out, fc_num));

				return PyObject_CallFunction(decimal_constructor, "s", fc_str);
			}

		case jbvString:
			return fc_PLyUnicode_FromJsonbValue(fc_jsonbValue);

		case jbvBool:
			if (fc_jsonbValue->val.boolean)
				Py_RETURN_TRUE;
			else
				Py_RETURN_FALSE;

		default:
			elog(ERROR, "unexpected jsonb value type: %d", fc_jsonbValue->type);
			return NULL;
	}
}

/*
 * PLyObject_FromJsonbContainer
 *
 * 将 JsonbContainer 转换为 PyObject。
 */
static PyObject * fc_PLyObject_FromJsonbContainer(JsonbContainer *fc_jsonb)
{
	JsonbIteratorToken fc_r;
	JsonbValue	fc_v;
	JsonbIterator *fc_it;
	PyObject   *fc_result;

	fc_it = JsonbIteratorInit(fc_jsonb);
	fc_r = JsonbIteratorNext(&fc_it, &fc_v, true);

	switch (fc_r)
	{
		case WJB_BEGIN_ARRAY:
			if (fc_v.val.array.rawScalar)
			{
				JsonbValue	fc_tmp;

				if ((fc_r = JsonbIteratorNext(&fc_it, &fc_v, true)) != WJB_ELEM ||
					(fc_r = JsonbIteratorNext(&fc_it, &fc_tmp, true)) != WJB_END_ARRAY ||
					(fc_r = JsonbIteratorNext(&fc_it, &fc_tmp, true)) != WJB_DONE)
					elog(ERROR, "unexpected jsonb token: %d", fc_r);

				fc_result = fc_PLyObject_FromJsonbValue(&fc_v);
			}
			else
			{
				PyObject   *volatile fc_elem = NULL;

				fc_result = PyList_New(0);
				if (!fc_result)
					return NULL;

				PG_TRY();
				{
					while ((fc_r = JsonbIteratorNext(&fc_it, &fc_v, true)) != WJB_DONE)
					{
						if (fc_r != WJB_ELEM)
							continue;

						fc_elem = fc_PLyObject_FromJsonbValue(&fc_v);

						PyList_Append(fc_result, fc_elem);
						Py_XDECREF(fc_elem);
						fc_elem = NULL;
					}
				}
				PG_CATCH();
				{
					Py_XDECREF(fc_elem);
					Py_XDECREF(fc_result);
					PG_RE_THROW();
				}
				PG_END_TRY();
			}
			break;

		case WJB_BEGIN_OBJECT:
			{
				PyObject   *volatile fc_result_v = PyDict_New();
				PyObject   *volatile fc_key = NULL;
				PyObject   *volatile fc_val = NULL;

				if (!fc_result_v)
					return NULL;

				PG_TRY();
				{
					while ((fc_r = JsonbIteratorNext(&fc_it, &fc_v, true)) != WJB_DONE)
					{
						if (fc_r != WJB_KEY)
							continue;

						fc_key = fc_PLyUnicode_FromJsonbValue(&fc_v);
						if (!fc_key)
						{
							Py_XDECREF(fc_result_v);
							fc_result_v = NULL;
							break;
						}

						if ((fc_r = JsonbIteratorNext(&fc_it, &fc_v, true)) != WJB_VALUE)
							elog(ERROR, "unexpected jsonb token: %d", fc_r);

						fc_val = fc_PLyObject_FromJsonbValue(&fc_v);
						if (!fc_val)
						{
							Py_XDECREF(fc_key);
							fc_key = NULL;
							Py_XDECREF(fc_result_v);
							fc_result_v = NULL;
							break;
						}

						PyDict_SetItem(fc_result_v, fc_key, fc_val);

						Py_XDECREF(fc_key);
						fc_key = NULL;
						Py_XDECREF(fc_val);
						fc_val = NULL;
					}
				}
				PG_CATCH();
				{
					Py_XDECREF(fc_result_v);
					Py_XDECREF(fc_key);
					Py_XDECREF(fc_val);
					PG_RE_THROW();
				}
				PG_END_TRY();

				fc_result = fc_result_v;
			}
			break;

		default:
			elog(ERROR, "unexpected jsonb token: %d", fc_r);
			return NULL;
	}

	return fc_result;
}

/*
 * PLyMapping_ToJsonbValue
 *
 * 将 Python 字典转换为 JsonbValue。
 */
static JsonbValue * fc_PLyMapping_ToJsonbValue(PyObject *fc_obj, JsonbParseState **fc_jsonb_state)
{
	Py_ssize_t	fc_pcount;
	PyObject   *volatile fc_items;
	JsonbValue *volatile fc_out;

	fc_pcount = PyMapping_Size(fc_obj);
	fc_items = PyMapping_Items(fc_obj);

	PG_TRY();
	{
		Py_ssize_t	fc_i;

		pushJsonbValue(fc_jsonb_state, WJB_BEGIN_OBJECT, NULL);

		for (fc_i = 0; fc_i < fc_pcount; fc_i++)
		{
			JsonbValue	fc_jbvKey;
			PyObject   *fc_item = PyList_GetItem(fc_items, fc_i);
			PyObject   *fc_key = PyTuple_GetItem(fc_item, 0);
			PyObject   *fc_value = PyTuple_GetItem(fc_item, 1);

			/* Python 字典可以有 None 作为键 */
			if (fc_key == Py_None)
			{
				fc_jbvKey.type = jbvString;
				fc_jbvKey.val.string.len = 0;
				fc_jbvKey.val.string.val = "";
			}
			else
			{
				/* 所有其他类型的键我们序列化为字符串 */
				fc_PLyUnicode_ToJsonbValue(fc_key, &fc_jbvKey);
			}

			(void) pushJsonbValue(fc_jsonb_state, WJB_KEY, &fc_jbvKey);
			(void) fc_PLyObject_ToJsonbValue(fc_value, fc_jsonb_state, false);
		}

		fc_out = pushJsonbValue(fc_jsonb_state, WJB_END_OBJECT, NULL);
	}
	PG_FINALLY();
	{
		Py_DECREF(fc_items);
	}
	PG_END_TRY();

	return fc_out;
}

/*
 * PLySequence_ToJsonbValue
 *
 * 将 Python 列表转换为 JsonbValue。期望转换后的 PyObject 和
 * jsonb 构造所需的状态。
 */
static JsonbValue * fc_PLySequence_ToJsonbValue(PyObject *fc_obj, JsonbParseState **fc_jsonb_state)
{
	Py_ssize_t	fc_i;
	Py_ssize_t	fc_pcount;
	PyObject   *volatile fc_value = NULL;

	fc_pcount = PySequence_Size(fc_obj);

	pushJsonbValue(fc_jsonb_state, WJB_BEGIN_ARRAY, NULL);

	PG_TRY();
	{
		for (fc_i = 0; fc_i < fc_pcount; fc_i++)
		{
			fc_value = PySequence_GetItem(fc_obj, fc_i);
			Assert(fc_value);

			(void) fc_PLyObject_ToJsonbValue(fc_value, fc_jsonb_state, true);
			Py_XDECREF(fc_value);
			fc_value = NULL;
		}
	}
	PG_CATCH();
	{
		Py_XDECREF(fc_value);
		PG_RE_THROW();
	}
	PG_END_TRY();

	return pushJsonbValue(fc_jsonb_state, WJB_END_ARRAY, NULL);
}

/*
 * PLyNumber_ToJsonbValue(PyObject *obj)
 *
 * 将 Python 数字转换为 JsonbValue。
 */
static JsonbValue * fc_PLyNumber_ToJsonbValue(PyObject *fc_obj, JsonbValue *fc_jbvNum)
{
	Numeric		fc_num;
	char	   *fc_str = PLyObject_AsString(fc_obj);

	PG_TRY();
	{
		Datum		fc_numd;

		fc_numd = DirectFunctionCall3(numeric_in,
								   CStringGetDatum(fc_str),
								   ObjectIdGetDatum(InvalidOid),
								   Int32GetDatum(-1));
		fc_num = DatumGetNumeric(fc_numd);
	}
	PG_CATCH();
	{
		ereport(ERROR,
				(errcode(ERRCODE_DATATYPE_MISMATCH),
				 errmsg("could not convert value \"%s\" to jsonb", fc_str)));
	}
	PG_END_TRY();

	pfree(fc_str);

	/*
	 * jsonb 不允许 NaN 或无穷大（根据 JSON 规范），因此我们
	 * 必须在这里明确拒绝这些。
	 */
	if (numeric_is_nan(fc_num))
		ereport(ERROR,
				(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
				 errmsg("cannot convert NaN to jsonb")));
	if (numeric_is_inf(fc_num))
		ereport(ERROR,
				(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
				 errmsg("cannot convert infinity to jsonb")));

	fc_jbvNum->type = jbvNumeric;
	fc_jbvNum->val.numeric = fc_num;

	return fc_jbvNum;
}

/*
 * PLyObject_ToJsonbValue(PyObject *obj)
 *
 * 将 Python 对象转换为 JsonbValue。
 */
static JsonbValue * fc_PLyObject_ToJsonbValue(PyObject *fc_obj, JsonbParseState **fc_jsonb_state, bool fc_is_elem)
{
	JsonbValue *fc_out;

	if (!PyUnicode_Check(fc_obj))
	{
		if (PySequence_Check(fc_obj))
			return fc_PLySequence_ToJsonbValue(fc_obj, fc_jsonb_state);
		else if (PyMapping_Check(fc_obj))
			return fc_PLyMapping_ToJsonbValue(fc_obj, fc_jsonb_state);
	}

	fc_out = palloc(sizeof(JsonbValue));

	if (fc_obj == Py_None)
		fc_out->type = jbvNull;
	else if (PyUnicode_Check(fc_obj))
		fc_PLyUnicode_ToJsonbValue(fc_obj, fc_out);

	/*
	 * PyNumber_Check() 对布尔值返回 true，因此布尔检查应该
	 * 首先进行。
	 */
	else if (PyBool_Check(fc_obj))
	{
		fc_out->type = jbvBool;
		fc_out->val.boolean = (fc_obj == Py_True);
	}
	else if (PyNumber_Check(fc_obj))
		fc_out = fc_PLyNumber_ToJsonbValue(fc_obj, fc_out);
	else
		ereport(ERROR,
				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
				 errmsg("Python type \"%s\" cannot be transformed to jsonb",
						PLyObject_AsString((PyObject *) fc_obj->ob_type))));

	/* 将结果推入 'jsonb_state'，除非它是原始标量值。 */
	return (*fc_jsonb_state ?
			pushJsonbValue(fc_jsonb_state, fc_is_elem ? WJB_ELEM : WJB_VALUE, fc_out) :
			fc_out);
}

/*
 * plpython_to_jsonb
 *
 * 将 Python 对象转换为 Jsonb 数据
 */
PG_FUNCTION_INFO_V1(plpython_to_jsonb);
Datum plpython_to_jsonb(PG_FUNCTION_ARGS)
{
	PyObject   *fc_obj;
	JsonbValue *fc_out;
	JsonbParseState *fc_jsonb_state = NULL;

	fc_obj = (PyObject *) PG_GETARG_POINTER(0);
	fc_out = fc_PLyObject_ToJsonbValue(fc_obj, &fc_jsonb_state, true);
	PG_RETURN_POINTER(JsonbValueToJsonb(fc_out));
}

/*
 * jsonb_to_plpython
 *
 * 将 Jsonb 数据转换为 PyObject，并作为内部返回。
 */
PG_FUNCTION_INFO_V1(jsonb_to_plpython);
Datum jsonb_to_plpython(PG_FUNCTION_ARGS)
{
	PyObject   *fc_result;
	Jsonb	   *fc_in = PG_GETARG_JSONB_P(0);

	/*
	 * 初始化指向 Decimal 构造函数的指针。首先我们尝试 "cdecimal"，C
	 * 版本的 decimal 库。如果失败，我们使用较慢的 "decimal"
	 * 模块。
	 */
	if (!decimal_constructor)
	{
		PyObject   *fc_decimal_module = PyImport_ImportModule("cdecimal");

		if (!fc_decimal_module)
		{
			PyErr_Clear();
			fc_decimal_module = PyImport_ImportModule("decimal");
		}
		Assert(fc_decimal_module);
		decimal_constructor = PyObject_GetAttrString(fc_decimal_module, "Decimal");
	}

	fc_result = fc_PLyObject_FromJsonbContainer(&fc_in->root);
	if (!fc_result)
		PLy_elog(ERROR, "transformation from jsonb to Python failed");

	return PointerGetDatum(fc_result);
}
